import json
import math
import time
# 多线程
from threading import Thread

import cv2
import gym
import numpy as np
import win32con
import win32gui
import win32ui
from pykeyboard import PyKeyboard
from redis import Redis


# from pymouse import PyMouse


class HKEnv2D(gym.Env):
    def __init__(self):
        self.r = Redis(host='127.0.0.1', port=6379, db=1)
        self.hwnd_id = []
        self.k = PyKeyboard()
        self.initailize()
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)
        # 创建dc, 截取桌面
        self.hdesktop = win32gui.GetDesktopWindow()
        # 获取窗口大小
        l, t, r, b = win32gui.GetWindowRect(self.hwnd_id[0])
        self.width = r - l
        self.height = b - t

        # 创建设备描述表
        self.desktop_dc = win32gui.GetWindowDC(self.hdesktop)
        # 创建dc对象
        self.mfc_dc = win32ui.CreateDCFromHandle(self.desktop_dc)
        # 创建内存设备描述表
        self.save_dc = self.mfc_dc.CreateCompatibleDC()
        # 创建位图对象
        self.save_bit_map = win32ui.CreateBitmap()
        # 截图
        self.save_bit_map.CreateCompatibleBitmap(self.mfc_dc, self.width, self.height)

        # self.action_list = ['攻击', '左+攻', '右+攻', '上+攻', '下+攻', '左+跳', '右+跳', '左+跳+攻', '右+跳+攻', '波', '吼', '砸', '冲']
        # self.action_dict = {
        #     0: '攻击',
        #     1: '左+攻',
        #     2: '右+攻',
        #     3: '上+攻',
        #     4: '下+攻',
        #     5: '左+跳',
        #     6: '右+跳',
        #     7: '左+跳+攻',
        #     8: '右+跳+攻',
        #     9: '波',
        #     10: '吼',
        #     11: '砸',
        #     12: '冲'
        # }
        self._max_episode_steps = 10000
        self.observation_space = gym.spaces.Box(low=0, high=1, shape=(300, 300), dtype=np.float32)
        self.action_space = gym.spaces.Discrete(12)
        self.std = 10
        self.mean = -1
        self.steps = 2

        self.min_x = 10000
        self.max_x = 0

    def initailize(self):
        pass

    def get_all_hwnd(self, hwnd, mouse):
        if (win32gui.IsWindow(hwnd) and
            win32gui.IsWindowEnabled(hwnd) and
            win32gui.IsWindowVisible(hwnd)) and win32gui.GetWindowText(hwnd).__contains__('Hollow Knight'):
            self.hwnd_id.append(hwnd)

    def restart_game(self):
        win32gui.EnumWindows(self.get_all_hwnd, 0)
        time.sleep(5)
        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        time.sleep(1)
        self.k.release_key(self.k.left_key)
        self.k.release_key(self.k.right_key)
        self.k.release_key(self.k.up_key)
        self.k.release_key(self.k.down_key)
        self.k.release_key("Z")
        self.k.release_key("X")
        self.k.release_key("C")
        self.k.release_key("F")
        time.sleep(1)
        self.k.tap_key(self.k.left_key)
        time.sleep(4)
        self.k.tap_key(self.k.up_key)
        time.sleep(2)
        self.k.tap_key("z")

    def reset(self, *args, **kwargs):
        start = None
        binary = None
        while start is None:
            self.face = 1
            self._last_enemy_hp = 0
            self._last_knight_hp = 0
            self.can_shadow_dash = True
            self.shadow_dashing_time = 0
            self.s = time.time()
            self.e = time.time()
            t = Thread(target=self.restart_game)
            t.start()
            start = self.r.brpop(['start'], timeout=20)
            time.sleep(0.1)
            state_json = self.r.get('hollow_obs')
            binary = self.get_binary(save=True, no=0)
            self.state_to_obs(state_json)
            self._last_enemy_hp = self.enemy_hp
            self._last_knight_hp = self.knight_hp
        print("reset done")
        return binary

    def do_action(self, action):
        # self.action_dict = {0: '攻击', 1: '上+攻', 2: '下+攻', 3: '跳', 4: '松+跳(双跳)',
        #                     5: '波', 6: '吼', 7: '砸', 8: '冲', 9: '左', 10: '右', 11: '松(无动作)'}
        L = self.k.left_key
        R = self.k.right_key
        U = self.k.up_key
        D = self.k.down_key
        if action == 0:
            self.k.tap_key('X')
        if action == 1:
            self.k.release_key(D)
            self.k.press_key(U)
            time.sleep(0.01)
            self.k.tap_key('X')
            time.sleep(0.01)
            self.k.release_key(U)
        if action == 2:
            self.k.release_key(U)
            self.k.press_key(D)
            time.sleep(0.01)
            self.k.tap_key('X')
            time.sleep(0.01)
            self.k.release_key(D)
        if action == 3:
            self.k.press_key('Z')
        if action == 4:
            self.k.release_key('Z')
            time.sleep(0.01)
            self.k.press_key('Z')
        if action == 5:
            self.k.tap_key('F')
        if action == 6:
            self.k.release_key(D)
            self.k.press_key(U)
            time.sleep(0.01)
            self.k.tap_key('F')
            time.sleep(0.01)
            self.k.release_key(U)
        if action == 7:
            self.k.release_key(U)
            self.k.press_key(D)
            time.sleep(0.01)
            self.k.tap_key('F')
            time.sleep(0.01)
            self.k.release_key(D)
        if action == 8:
            self.k.tap_key("C")
        if action == 9:
            self.k.release_key(R)
            self.k.press_key(L)
        if action == 10:
            self.k.release_key(L)
            self.k.press_key(R)
        if action == 11:
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.release_key('Z')

    def get_binary(self, save=False, no=0):
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)

        self.save_dc.SelectObject(self.save_bit_map)
        # 截取从左上角(50, 50)长宽为(width, height)的图片
        # self.save_dc.BitBlt((0, 0), (self.width - 21, self.height - 41), self.mfc_dc, (0, 0), win32con.SRCCOPY)
        self.save_dc.BitBlt((0, 0), (self.width - 300, self.height - 160), self.mfc_dc, (150, 50), win32con.SRCCOPY)

        # 保存位图文件
        signedIntsArray = self.save_bit_map.GetBitmapBits(True)
        img = np.frombuffer(signedIntsArray, dtype='uint8')
        img.shape = (self.height, self.width, 4)
        # 删除多余部分
        img = img[0:self.height - 160, 0:self.width - 300]
        # 缩放
        img = cv2.resize(img, (300, 300))
        if save:
            cv2.imwrite("imgs/{}.png".format(no), img)
        # 保存
        # cv2.imwrite('test.png', img)
        # 转换为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if save:
            cv2.imwrite("imgs/gray_{}.png".format(no), gray)
        # 压缩到0-1
        gray = gray / 255.0
        # 缩放
        # gray = cv2.resize(gray, (511, 511))
        # 保存
        # cv2.imwrite('test_gray.png', gray)
        # 二值化
        # ret, binary = cv2.threshold(gray, 127, 1, cv2.THRESH_BINARY)
        # if save:
        #     ret, binary_255 = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        #     cv2.imwrite("imgs/binary_{}.png".format(no), binary_255)
        # 转成numpy数组
        binary = np.array(gray)
        return binary

    def step(self, action):
        self.steps += 1
        reward = -1
        if action == 11:
            reward -= 1
        if abs(self.enemy_x - self.knight_x) > 3:
            if self.enemy_x > self.knight_x:
                # 向左走
                if action == 9:
                    reward -= 1
                # 左冲
                if self.face == 0.5 and action == 8:
                    reward -= 1
            else:
                # 向右走
                if action == 10:
                    reward -= 1
                # 右冲
                if self.face == 1 and action == 8:
                    reward -= 1
        if action in [0, 1, 2]:
            reward -= 1
        before_mp = self.knight_mp
        start = time.time()
        self.do_action(action)
        max_time = 0.04
        while True:
            end = time.time()
            state_json = self.r.get('hollow_obs')
            self.state_to_obs(state_json)
            if str(self.r.get('done').decode('utf-8')) == "1":
                print('kill')
                break
            if self.knight_spell or self.dashing or self.is_attacking or end - start < max_time:
                time.sleep(0.001)
                continue
            if end - start < max_time + 0.01:
                continue
            else:
                break
        binary = self.get_binary()
        state_json = self.r.get('hollow_obs')
        self.state_to_obs(state_json)
        self.done = False
        self.key_point = False
        if self.distance > 6:
            reward -= self.distance
        if self.distance < 3:
            reward -= 2
        # if self.steps < 1000:
        #     if action in [5, 6, 7]:
        #         reward = 10
        if self._last_enemy_hp > self.enemy_hp:
            if action in [0, 1, 2]:
                reward = self.distance * 0.1
            if before_mp > 33:
                if action in [5, 6]:
                    reward = self.distance * 0.4
                if action == 7:
                    reward = self.distance * 0.2
            if self.distance > 6 and action in [5, 6, 7]:
                self.key_point = True
        if self._last_knight_hp > self.knight_hp:
            reward = -100
            self.key_point = True
        if self.knight_hp <= 0 or (self._last_knight_hp < 5 and self.knight_hp > 6):
            self.done = True
            reward = -100
            self.key_point = True
        if str(self.r.get('done').decode('utf-8')) == "1":
            try:
                self.get_binary(save=True, no=1)
            except:
                pass
            self.done = True
            if action in [0, 1, 2]:
                reward = min(self.distance * 0.1, 1)
            if before_mp > 33:
                if action in [5, 6]:
                    reward = min(self.distance * 0.4, 1)
                if action == 7:
                    reward = min(self.distance * 0.2, 1)
                if self.distance > 5:
                    self.key_point = True
        self._last_enemy_hp = self.enemy_hp
        self._last_knight_hp = self.knight_hp
        if self.done:
            self.k.release_key(self.k.left_key)
            self.k.release_key(self.k.right_key)
            self.k.release_key(self.k.up_key)
            self.k.release_key(self.k.down_key)
            self.k.release_key("Z")
            self.k.release_key("X")
            self.k.release_key("C")
            self.k.release_key("F")
        # 缩放reward
        # self.mean, self.std = self.calculate_mean_std(reward, self.mean, self.std)
        # reward = reward / self.std
        return binary, reward, self.done, self.key_point

    def calculate_mean_std(self, new_history, old_mean, old_std):
        # sample_count：旧样本总数
        # new_sample_num：新样本总数
        # new_history：新的样本
        # old_mean：旧均值
        # old_var：旧方差
        new_mean = (self.steps * old_mean + new_history)
        new_mean /= (1 + self.steps)
        new_std = np.sum(np.square(new_history - new_mean), axis=0)
        new_std = (self.steps * np.square(old_std) + new_std)
        new_std /= (1 + self.steps)
        new_std = np.sqrt(new_std)
        return new_mean, new_std

    def state_to_obs(self, state):
        state_dict = json.loads(state)
        # 是否黑冲
        self.shadow_dashing = True if state_dict['shadow_dashing'] == '1' else False

        # 是否冲刺
        self.dashing = True if state_dict['dashing'] == '1' else False

        # 是否在攻击
        self.is_attacking = True if state_dict['is_attacking'] == '1' else False

        # 是否在释放技能
        self.knight_spell = True if state_dict['knight_spell'] != "Inactive" else False

        # 是否在地面
        self.on_ground = True if state_dict['onGround'] == '1' else False
        # 血量
        self.knight_hp = int(state_dict['knight_hp'])
        self.enemy_hp = int(state_dict['enemy_hp'])

        # knight 的位置
        self.knight_x = float(state_dict['knight_x'])
        self.knight_y = float(state_dict['knight_y'])

        # knight的灵魂量
        self.knight_mp = float(state_dict['knight_mp'])

        # boss 的位置
        self.enemy_x = float(state_dict['enemy_x'])
        self.enemy_y = float(state_dict['enemy_y'])

        # 距离, 用于判断是否在攻击范围内, 平方和开方,
        self.distance = math.sqrt((self.knight_x - self.enemy_x) ** 2 + (self.knight_y - self.enemy_y) ** 2)

        # 冲刺是否冷却
        self.is_sprint = True

        # 相对位置
        self.relative_x = self.knight_x - self.enemy_x
        self.relative_y = self.knight_y - self.enemy_y

        self.angle_kb = self.angle(-self.relative_x, -self.relative_y)

    def angle(self, x, y):
        a = np.arctan2(y, x)
        if a < 0:
            a = a + 2 * np.pi
        return a


if __name__ == '__main__':
    env = HKEnv2D()
    # env.restart_game()z
    win32gui.EnumWindows(env.get_all_hwnd, 0)
    handle = env.hwnd_id[0]
    # 激活窗口
    win32gui.SetForegroundWindow(handle)
    time.sleep(3)
    k = PyKeyboard()
    k.press_key("Z")
    time.sleep(0.2)
    k.release_key("Z")
    time.sleep(0.1)
    k.press_key("Z")
    time.sleep(1)
    k.release_key('Z')

    obs = env.reset()
    print(obs)
    for i in range(100):
        action = np.random.randint(0, 12)
        obs, reward, done, info = env.step(1, 1)
        count = 0
        print(obs, reward, done, info)
